from typing import Optional
from .translation_chain import TranslationChain
from .pdf_parser import PDFParser
from .writer import Writer
from ..utils import LOG


class PDFTranslator:
    def __init__(self, model_name: str):
        self.translatoe_chain = TranslationChain(model_name)
        self.pdf_parser = PDFParser()
        self.writer = Writer()

    def translate_pdf(self, input_file: str, output_file_format: str = "markdown", source_language: str = "English", target_language: str = "Chinese", start_page: int = 1,
                      end_page: Optional[int] = None):
        self.book = self.pdf_parser.parse_pdf(input_file, start_page, end_page)
        LOG.debug(f"PDF 文件解析完成，共 {len(self.book.pages)} 页。")

        for page_idx, page in enumerate(self.book.pages):
            LOG.info(f"正在翻译第 {page_idx + 1} 页...")
            for content_idx, content in enumerate(page.contents):
                translation, status = self.translatoe_chain.run(content, source_language, target_language)
                self.book.pages[page_idx].contents[content_idx].set_translation(translation, status)

        return self.writer.save_translated_book(self.book, output_file_format)
